#!/usr/bin/env python
# encoding: utf-8
import pickle
import mysql.connector
from general.mysqlconfig import *
from general.generaltools import Debugger
class MySQL(Debugger):
    """This class describes all the packed SQL operations required by entity-
    linking.
    """
    def __init__(self, fetchsize=1000000, isDebug=False, mode='std', use_db=False):
        Debugger.__init__(self, isDebug, mode)
        self.fetchsize = fetchsize
        self.connection = ""
        if use_db:
            self.init_db()

    def init_db(self):
        self.kpp_db = pickle.load(open('./general/db/kpp.db', 'rb'))
        self.entity_db = pickle.load(open('./general/db/entity.db', 'rb'))
        self.category_db = pickle.load(open('./general/db/category.db', 'rb'))

    def connect(self, user=user, password=password, database=database):
        try:
            self.connection = mysql.connector.connect(user=user, \
                    password=password, database=database)
        except:
            self.bugInfo += "Error: Can't not connect to MySQL\n\n"

    def release(self):
        if not self.connection:
            return
        self.connection.close()

    def get_keyphraseness(self, *ids):
        if not self.connection:
            return
        #where_condition = "content=\"" + "\" or content=\"".join(ids) + "\""
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = """SELECT id, content, KeywordPercent FROM hz_lexicon WHERE %s;""" % (where_condition)
        self.dPrint(select_sql)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                yield row[0], row[1], row[2]
        cursor.close()

    def get_kpp(self, *ids):
        if not self.connection:
            return
        if not len(ids):
            #print('Error: No ids from get_kpp')
            return
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = "SELECT id, entityList, PercentList FROM hz_kpp WHERE {0};".format(where_condition)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(int, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [int(row[2])]
                yield row[0], row1, row2
        cursor.close()

    def build_kpp_mem(self):
        select_sql = "SELECT id, entityList, PercentList FROM hz_kpp;"
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        kpp_db = {}
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(float, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [float(row[2])]
                kpp_db[row[0]] = [row1, row2]
        mydb = open("kpp.db", 'wb')
        pickle.dump(kpp_db, mydb)
        mydb.close()

    def get_kpp_from_mem(self, *ids):
        if not len(ids):
            #print("Error: No ids from get kpp from mem")
            return
        ids = list(map(int, ids))
        for keyword_id in ids:
            if keyword_id in self.kpp_db:
                info = self.kpp_db[keyword_id]
                yield keyword_id, info[0], info[1]

    def get_entity_info(self, *ids):
        if not self.connection:
            return
        if not len(ids):
            #print('Error: No ids from get entity info')
            return
        ids = list(map(str, ids))
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = "SELECT Id, InlinkIdList, InlinkFreqList, " + \
                "OutlinkIdList, OutlinkFreqList FROM hz_entity WHERE {0};".format(where_condition)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(int, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [int(row[2])]
                if not row[3]:
                    row3 = []
                else:
                    row3 = list(map(int, row[3].split(',')[:-1])) if row[3].find(',') != -1 else [int(row[3])]
                if not row[4]:
                    row4 = []
                else:
                    row4 = list(map(int, row[4].split(',')[:-1])) if row[4].find(',') != -1 else [int(row[4])]
                yield row[0], row1, row2, row3, row4
        cursor.close()

    def build_entity_mem(self):
        select_sql = "SELECT Id, InlinkIdList, InlinkFreqList, " + \
                "OutlinkIdList, OutlinkFreqList, CategoryIdList FROM hz_entity;"
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        entity_db = {}
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(int, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [int(row[2])]
                if not row[3]:
                    row3 = []
                else:
                    row3 = list(map(int, row[3].split(',')[:-1])) if row[3].find(',') != -1 else [int(row[3])]
                if not row[4]:
                    row4 = []
                else:
                    row4 = list(map(int, row[4].split(',')[:-1])) if row[4].find(',') != -1 else [int(row[4])]
                if not row[5]:
                    row5 = []
                else:
                    row5 = list(map(int, row[5].split(',')[:-1])) if row[5].find(',') != -1 else [int(row[5])]
                entity_db[row[0]] = [row1, row2, row3, row4, row5]
        mydb = open("entity.db", 'wb')
        pickle.dump(entity_db, mydb)
        mydb.close()

    def get_entity_info_from_mem(self, *ids):
        if not len(ids):
            #print("Error: No ids from get entity info from mem")
            return
        ids = list(map(int, ids))
        for entity_id in ids:
            entity_info = self.entity_db[entity_id]
            yield entity_id, entity_info[0], entity_info[1], \
                    entity_info[2], entity_info[3], entity_info[4]

    def get_entity_link_info_from_mem(self, *ids):
        if not len(ids):
            #print("Error: No ids from get entity info from mem")
            return
        ids = list(map(int, ids))
        for entity_id in ids:
            entity_info = self.entity_db[entity_id]
            yield entity_id, entity_info[0], entity_info[1], entity_info[2], entity_info[3]

    def get_second_order_inlink(self, *ids):
        if not len(ids):
            return []
        second_order_inlink = set()
        for entity_id in ids:
            if entity_id not in self.entity_db: continue #Need deal
            second_order_inlink.update(self.entity_db[entity_id][0])
        return list(second_order_inlink)

    def get_entity_category_info_from_mem(self, *ids):
        if not len(ids):
            #print("Error: No ids from get entity category info")
            return
        for entity_id in ids:
            entity_info = self.entity_db[entity_id]
            yield entity_id, entity_info[4]

    def get_links_freq(self, *ids):
        if not self.connection:
            return
        if not len(ids):
            #print('Error: No ids from get links freq')
            return
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = "SELECT Id, OutlinkIdList, OutlinkFreqList, " + \
                "InlinkIdList, InlinkFreqList FROM hz_entity WHERE {0};".format(where_condition)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(int, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [int(row[2])]
                if not row[3]:
                    row3 = []
                else:
                    row3 = list(map(int, row[3].split(',')[:-1])) if row[3].find(',') != -1 else [int(row[3])]
                if not row[4]:
                    row4 = []
                else:
                    row4 = list(map(int, row[4].split(',')[:-1])) if row[4].find(',') != -1 else [int(row[4])]
                yield row[0], row1, row2, row3, row4
        cursor.close()

    def get_idf(self, *ids):
        if not self.connection:
            return
        if not len(ids):
            #print('Error: No ids from get idf')
            return
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = "SELECT Id, idf FROM hz_dictionary WHERE {0};".format(where_condition)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                yield row[0], row[1]
        cursor.close()

    def get_tf_dict(self, *ids):
        if not self.connection:
            return
        if not len(ids):
            #print('Error: No ids from get tf dict')
            return
        if len(ids) == 1:
            where_condition = "id={0}".format(ids[0])
        else:
            where_condition = "id=" + " or id=".join(ids)
        select_sql = "SELECT Id, word_id_list, tf_list, fl FROM hz_entity WHERE {0};".format(where_condition)
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                if not row[1]:
                    row1 = []
                else:
                    row1 = list(map(int, row[1].split(',')[:-1])) if row[1].find(',') != -1 else [int(row[1])]
                if not row[2]:
                    row2 = []
                else:
                    row2 = list(map(int, row[2].split(',')[:-1])) if row[2].find(',') != -1 else [int(row[2])]
                yield row[0], dict(zip(row1, row2)), row[3]
        cursor.close()

    def build_category_mem(self):
        select_sql = "SELECT cateId, cateFreq FROM hz_category;"
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        category_db = {}
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                category_db[row[0]] = row[1]
        mydb = open("category.db", 'wb')
        pickle.dump(category_db, mydb)
        mydb.close()

    def get_category_info_from_mem(self, *ids):
        if not len(ids):
            #print("Error: No ids from get category info from mem")
            return
        ids = list(map(int, ids))
        for category_id in ids:
            if category_id not in self.category_db:
                print("Error: category not fouond:", category_id, type(category_id))
                continue
            yield category_id, self.category_db[category_id]

    def get_all_keyword(self):
        select_sql = "SELECT id, content FROM hz_lexicon;"
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                yield row[0], row[1]
        cursor.close()

    def get_all_entity_id(self):
        select_sql = "SELECT Id, Title FROM hz_entity;"
        cursor = self.connection.cursor()
        cursor.execute(select_sql)
        while True:
            rows = cursor.fetchmany(self.fetchsize)
            if not rows: break
            for row in rows:
                yield row[0], row[1]
        cursor.close()

if __name__ == "__main__":
    mysql = MySQL()
    mysql.connect()
    mysql.build_kpp_mem()
    mysql.release()


#insert_sql = """INSERT INTO zs(name, age) VALUES (%s, %s)"""
#cursor.execute(insert_sql, ('Jay', 22))
#cursor.execute(insert_sql, ('张晟', 24))
#
#connection.commit()
#cursor.close()
